import tkinter as tk
from tkinter import messagebox
import os
import sys
from environment import WumpusWorld, Action
from agent import QLearningAgent

class WumpusWorldGUI:
    """Wumpus世界图形界面"""
    
    def __init__(self, env, agent, model_path=None):
        """
        初始化图形界面
        
        参数:
        - env: Wumpus环境
        - agent: 强化学习代理
        - model_path: 模型路径，如果提供则加载模型
        """
        self.env = env
        self.agent = agent
        
        # 如果提供了模型路径，则加载模型
        if model_path and os.path.exists(model_path):
            self.agent.load_q_table(model_path)
        
        self.cell_size = 80
        self.setup_gui()
    
    def setup_gui(self):
        """设置GUI"""
        self.window = tk.Tk()
        self.window.title("Wumpus世界 - 强化学习")
        
        size = self.env.size
        
        # 顶部信息栏
        frame_info = tk.Frame(self.window)
        frame_info.pack(pady=10)
        
        self.label_score = tk.Label(frame_info, text="分数: 0", font=("微软雅黑", 12))
        self.label_score.pack(side=tk.LEFT, padx=20)
        
        self.label_status = tk.Label(frame_info, text="状态: 游戏开始", font=("微软雅黑", 12))
        self.label_status.pack(side=tk.LEFT, padx=20)
        
        # 创建画布
        canvas_width = self.cell_size * size
        canvas_height = self.cell_size * size
        self.canvas = tk.Canvas(self.window, width=canvas_width, height=canvas_height, bg="white")
        self.canvas.pack(pady=10)
        
        # 控制面板
        frame_control = tk.Frame(self.window)
        frame_control.pack(pady=10)
        
        self.btn_reset = tk.Button(frame_control, text="重置", font=("微软雅黑", 12),
                                  command=self.reset_game)
        self.btn_reset.pack(side=tk.LEFT, padx=10)
        
        self.btn_step = tk.Button(frame_control, text="单步执行", font=("微软雅黑", 12),
                                 command=self.step_game)
        self.btn_step.pack(side=tk.LEFT, padx=10)
        
        self.btn_auto = tk.Button(frame_control, text="自动执行", font=("微软雅黑", 12),
                                 command=self.auto_play)
        self.btn_auto.pack(side=tk.LEFT, padx=10)
        
        self.btn_quit = tk.Button(frame_control, text="退出", font=("微软雅黑", 12),
                                 command=self.window.quit)
        self.btn_quit.pack(side=tk.LEFT, padx=10)
        
        # 代理状态信息
        frame_agent = tk.Frame(self.window)
        frame_agent.pack(pady=10)
        
        self.label_agent = tk.Label(frame_agent, text="代理状态: 拥有箭", font=("微软雅黑", 12))
        self.label_agent.pack()
        
        # 感知信息
        frame_percepts = tk.Frame(self.window)
        frame_percepts.pack(pady=10)
        
        self.label_percepts = tk.Label(frame_percepts, text="感知: 无", font=("微软雅黑", 12))
        self.label_percepts.pack()
        
        # 初始化游戏
        self.reset_game()
        
        # 自动播放标志
        self.auto_playing = False
    
    def reset_game(self):
        """重置游戏"""
        self.env.reset()
        self.update_display()
        self.auto_playing = False
        if hasattr(self, 'btn_auto'):
            self.btn_auto.config(text="自动执行")
    
    def step_game(self):
        """执行一步游戏"""
        if self.env.game_over:
            messagebox.showinfo("游戏结束", "游戏已结束，请重置游戏")
            return
        
        # 获取代理的动作
        state = self.env._get_state()
        action = self.agent.get_action(state, training=False)
        
        # 执行动作
        next_state, reward, done, info = self.env.step(action)
        
        # 更新显示
        self.update_display()
        
        # 显示动作信息
        action_names = ["上移", "右移", "下移", "左移", "拾取", "向上射箭", "向右射箭", "向下射箭", "向左射箭", "爬出"]
        action_info = f"执行动作: {action_names[action]}"
        if info["message"]:
            action_info += f" - {info['message']}"
        print(action_info)
        
        # 检查游戏是否结束
        if done:
            if self.env.won:
                messagebox.showinfo("游戏结束", f"恭喜! 代理成功带着金子离开洞穴!\n最终分数: {self.env.score}")
            else:
                messagebox.showinfo("游戏结束", f"游戏结束!\n最终分数: {self.env.score}")
    
    def auto_play(self):
        """自动执行游戏"""
        if self.env.game_over:
            messagebox.showinfo("游戏结束", "游戏已结束，请重置游戏")
            return
        
        self.auto_playing = not self.auto_playing
        
        if self.auto_playing:
            self.btn_auto.config(text="停止自动")
            self.auto_step()
        else:
            self.btn_auto.config(text="自动执行")
    
    def auto_step(self):
        """自动执行的单步"""
        if not self.auto_playing or self.env.game_over:
            self.auto_playing = False
            self.btn_auto.config(text="自动执行")
            return
        
        # 执行一步
        self.step_game()
        
        # 如果游戏未结束，继续自动执行
        if not self.env.game_over:
            self.window.after(500, self.auto_step)
        else:
            self.auto_playing = False
            self.btn_auto.config(text="自动执行")
    
    def update_display(self):
        """更新显示"""
        self.canvas.delete("all")
        
        size = self.env.size
        cell_size = self.cell_size
        
        # 绘制网格
        for i in range(size+1):
            # 水平线
            self.canvas.create_line(0, i*cell_size, size*cell_size, i*cell_size, width=2)
            # 垂直线
            self.canvas.create_line(i*cell_size, 0, i*cell_size, size*cell_size, width=2)
        
        # 绘制陷阱
        for pit_pos in self.env.pit_positions:
            x, y = pit_pos
            self.canvas.create_text((y+0.5)*cell_size, (x+0.5)*cell_size, 
                                   text="陷阱", font=("微软雅黑", 14), fill="black")
        
        # 绘制Wumpus
        wx, wy = self.env.wumpus_pos
        if self.env.wumpus_alive:
            self.canvas.create_text((wy+0.5)*cell_size, (wx+0.5)*cell_size, 
                                   text="Wumpus", font=("微软雅黑", 14), fill="red")
        else:
            self.canvas.create_text((wy+0.5)*cell_size, (wx+0.5)*cell_size, 
                                   text="死亡Wumpus", font=("微软雅黑", 12), fill="gray")
        
        # 绘制金子
        if not self.env.has_gold:
            gx, gy = self.env.gold_pos
            self.canvas.create_text((gy+0.5)*cell_size, (gx+0.5)*cell_size, 
                                   text="金子", font=("微软雅黑", 14), fill="gold")
        
        # 绘制起点
        sx, sy = self.env.start_pos
        self.canvas.create_text((sy+0.5)*cell_size, (sx+0.5)*cell_size, 
                               text="起点", font=("微软雅黑", 14), fill="blue")
        
        # 绘制代理
        ax, ay = self.env.agent_pos
        agent_text = "代理"
        if self.env.has_gold:
            agent_text += "🔶"  # 金子符号
        self.canvas.create_oval((ay+0.2)*cell_size, (ax+0.2)*cell_size, 
                               (ay+0.8)*cell_size, (ax+0.8)*cell_size, 
                               fill="green", outline="black")
        self.canvas.create_text((ay+0.5)*cell_size, (ax+0.5)*cell_size, 
                               text=agent_text, font=("微软雅黑", 12), fill="white")
        
        # 更新分数和状态
        self.label_score.config(text=f"分数: {self.env.score}")
        
        if self.env.game_over:
            if self.env.won:
                self.label_status.config(text="状态: 胜利!")
            else:
                self.label_status.config(text="状态: 失败!")
        else:
            self.label_status.config(text="状态: 游戏进行中")
        
        # 更新代理状态
        agent_info = []
        if self.env.has_arrow:
            agent_info.append("拥有箭")
        if self.env.has_gold:
            agent_info.append("拥有金子")
        self.label_agent.config(text=f"代理状态: {', '.join(agent_info) if agent_info else '无'}")
        
        # 更新感知
        percepts = self.env.get_percepts()
        percept_str = []
        if percepts["breeze"]:
            percept_str.append("微风")
        if percepts["stench"]:
            percept_str.append("臭气")
        if percepts["glitter"]:
            percept_str.append("金光")
        if percepts["scream"]:
            percept_str.append("尖叫")
        self.label_percepts.config(text=f"感知: {', '.join(percept_str) if percept_str else '无'}")
    
    def run(self):
        """运行GUI"""
        self.window.mainloop()

def main():
    """GUI脚本的主函数"""
    # 检查参数
    output_dir = "output"
    model_path = os.path.join(output_dir, "wumpus_q_table_final.npy")
    
    if not os.path.exists(model_path):
        print(f"错误: 找不到模型文件 {model_path}")
        print("请先运行 train.py 训练模型")
        sys.exit(1)
    
    # 创建环境
    env = WumpusWorld(size=4, pit_prob=0.15)
    
    # 创建代理
    state_size = 4 * 4 * 2 * 2 * 2  # (x, y, 是否有金子, 是否有箭, Wumpus是否活着)
    action_size = len(Action)
    agent = QLearningAgent(state_size=state_size, action_size=action_size)
    
    # 加载模型
    if not agent.load_q_table(model_path):
        print("无法加载模型，将使用未训练的代理")
    
    # 创建并运行GUI
    print("启动Wumpus世界GUI...")
    gui = WumpusWorldGUI(env, agent, model_path)
    gui.run()

if __name__ == "__main__":
    main()